"""
数据库连接和会话管理 - 支持同步和异步
"""

from sqlalchemy import create_engine
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
from sqlalchemy.ext.declarative import declarative_base
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool
from typing import Generator, AsyncGenerator

from app.config import settings


# 同步数据库引擎 - 用于Celery任务和初始化脚本
# 跨地域优化配置：香港服务器连接国内MySQL
engine = create_engine(
    settings.DATABASE_URL,
    poolclass=QueuePool,
    pool_size=20,              # 增加连接池大小
    max_overflow=30,           # 增加溢出连接数
    pool_pre_ping=True,        # 连接前ping测试
    pool_recycle=1800,         # 30分钟回收连接（跨地域网络不稳定）
    pool_timeout=30,           # 获取连接超时30秒
    echo=settings.DEBUG,
    # MySQL跨地域连接优化配置
    connect_args={
        "charset": "utf8mb4",
        "use_unicode": True,
        "autocommit": False,
        "connect_timeout": 10,  # 连接超时10秒
        "read_timeout": 30,     # 读取超时30秒
        "write_timeout": 30,    # 写入超时30秒
        # 设置SQL模式
        "init_command": "SET sql_mode='STRICT_TRANS_TABLES,NO_ZERO_DATE,NO_ZERO_IN_DATE,ERROR_FOR_DIVISION_BY_ZERO'"
    }
)

# 异步数据库引擎 - 用于FastAPI异步接口
# 跨地域优化配置：香港服务器连接国内MySQL
async_database_url = settings.DATABASE_URL.replace("mysql+pymysql://", "mysql+aiomysql://")
async_engine = create_async_engine(
    async_database_url,
    pool_size=20,              # 增加异步连接池大小
    max_overflow=30,           # 增加溢出连接数
    pool_pre_ping=True,        # 连接前ping测试
    pool_recycle=1800,         # 30分钟回收连接
    pool_timeout=30,           # 获取连接超时30秒
    echo=settings.DEBUG,
    # MySQL异步跨地域连接配置
    connect_args={
        "charset": "utf8mb4",
        "use_unicode": True,
        "autocommit": False,
        "connect_timeout": 10,  # 连接超时10秒
        "read_timeout": 30,     # 读取超时30秒
        "write_timeout": 30,    # 写入超时30秒
    }
)

# 同步会话工厂
SessionLocal = sessionmaker(
    autocommit=False,
    autoflush=False,
    bind=engine
)

# 异步会话工厂
AsyncSessionLocal = sessionmaker(
    async_engine,
    class_=AsyncSession,
    expire_on_commit=False
)

# 创建基础模型类
Base = declarative_base()


def get_db() -> Generator[Session, None, None]:
    """
    获取同步数据库会话
    用于Celery任务和同步操作
    """
    db = SessionLocal()
    try:
        yield db
    finally:
        db.close()


async def get_async_db() -> AsyncGenerator[AsyncSession, None]:
    """
    获取异步数据库会话
    用于FastAPI异步接口
    """
    async with AsyncSessionLocal() as session:
        try:
            yield session
        finally:
            await session.close()


def get_async_session():
    """
    获取异步数据库会话上下文管理器
    用于Celery任务中的异步操作
    """
    return AsyncSessionLocal()


def create_tables():
    """创建所有数据库表"""
    Base.metadata.create_all(bind=engine)


def drop_tables():
    """删除所有数据库表"""
    Base.metadata.drop_all(bind=engine)
